package hku.cs.lunuo.instrument;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import soot.Body;
import soot.BodyTransformer;
import soot.Local;
import soot.Scene;
import soot.SootClass;
import soot.SootMethod;
import soot.Unit;
import soot.jimple.IdentityStmt;
import soot.jimple.IfStmt;
import soot.jimple.IntConstant;
import soot.jimple.InvokeExpr;
import soot.jimple.InvokeStmt;
import soot.jimple.Jimple;
import soot.jimple.LookupSwitchStmt;
import soot.jimple.StringConstant;
import soot.jimple.TableSwitchStmt;
import soot.toolkits.graph.Block;
import soot.toolkits.graph.BlockGraph;
import soot.toolkits.graph.ExceptionalBlockGraph;

public class PredChianInstrument extends BodyTransformer {

	private static PredChianInstrument instance = new PredChianInstrument();

	private PredChianInstrument() {
	}

	public static PredChianInstrument v() {
		return instance;
	}
	

	SootClass counterClass;
	SootMethod increaseES, flushES, increaseCounter, add, recordLast,
			confirmLast;

	@Override
	protected void internalTransform(Body b, String arg1, Map arg2) {
		counterClass = Scene.v()
				.getSootClass("hku.cs.lunuo.instrument.Counter");
		increaseES = counterClass
				.getMethod("void increaseES(java.lang.String,java.lang.String,int)");
		increaseCounter = counterClass
				.getMethod("void increase(java.lang.String,java.lang.String,int)");
		flushES = counterClass
				.getMethod("void flushES(java.lang.String,java.lang.String,int,int,int)");
		add = counterClass
				.getMethod("void add(java.lang.String,java.lang.String,java.lang.String)");

		BlockGraph bGraph = new ExceptionalBlockGraph(b);

		Iterator<Block> blockIt = bGraph.getBlocks().iterator();
		Local counterLocal = Jimple.v().newLocal("counterRef",
				counterClass.getType());
		b.getLocals().add(counterLocal);
		String methodName = b.getMethod().toString();
		String className = methodName.substring(1, methodName.indexOf(":"));
		methodName = methodName.substring(methodName.indexOf(":") + 2,
				methodName.length() - 1);

		List<Block> blockLs = bGraph.getBlocks();
		
		for (int i = 0; i < blockLs.size(); i++) {
			Unit ut = blockLs.get(i).getTail();
			if (ut instanceof IfStmt || ut instanceof TableSwitchStmt
					|| ut instanceof LookupSwitchStmt) {
				String startIs = "";
				List<Block> preBs = blockLs.get(i).getPreds();

				if (preBs.size() < 2)
					startIs = String.valueOf(i);
				else {
					for (Block preb : preBs) {
						startIs += "_"
								+ String.valueOf(preb.getIndexInMethod());
					}
					startIs = startIs.substring(1, startIs.length());
				}

				if (ut instanceof IfStmt) {
					int j = i;
					
					List<Block> preStartBls = blockLs.get(j).getPreds();
					boolean isStartBlkOk = true;
					for (Block bses : preStartBls)
						if ( bses.getIndexInMethod() < i || bses.getIndexInMethod() > j) {
							isStartBlkOk = false;
							break;
						}
					if(isStartBlkOk)
					
					
					while (j < blockLs.size()
							&& blockLs.get(j + 1).getTail() instanceof IfStmt) {
						List<Block> preBls = blockLs.get(j + 1).getPreds();
						boolean sb = false;
						for (Block bsss : preBls)
							if (bsss.getIndexInMethod() < i
									|| bsss.getIndexInMethod() > j) {
								sb = true;
								break;
							}

						if (!sb)
							j++;
						else
							break;

					}
					int k = i;
					while (k <= j) {
						if (k == j) {
							List args = new ArrayList();
							args.add(StringConstant.v(className));
							args.add(StringConstant.v(methodName));
							args.add(StringConstant.v(startIs));
							args.add(IntConstant.v(k));
							args.add(IntConstant.v(-(k + 1)));
							// here make negative value means false branch
							InvokeExpr incExpr1 = Jimple.v()
									.newStaticInvokeExpr(flushES.makeRef(),
											args);
							InvokeStmt incStmt1 = Jimple.v().newInvokeStmt(
									incExpr1);
							
							insertUnitToHead(incStmt1, blockLs.get(k + 1));
							
						}

						List<Block> sucBlks = blockLs.get(k).getSuccs();
						List<Integer> sucRightBlock = new ArrayList<Integer>();

						for (Block sucb : sucBlks) {

							if (!(sucb.getHead() instanceof IdentityStmt)) {
								int targetBlkIdx = sucb.getIndexInMethod();
								if (targetBlkIdx != k + 1) {
									sucRightBlock.add(targetBlkIdx);
								}
							}
						}

						int targetBlkIdx;

						if (sucRightBlock.size() > 1) {
							System.out.println("current block " + k);
							System.out.println("suc Block " + sucRightBlock);
							System.out.println(className + " " + methodName);
							throw new IllegalStateException(
									"must not have more than two suc normal block");
						} else if (sucRightBlock.size() == 0)
							targetBlkIdx = k + 1;// deal empty if block or while
						// block, may be this will
						// has no true branch
						else
							targetBlkIdx = sucRightBlock.get(0);

						// targetBlkIdx <= i because while structure will loop
						// in one block
						// Block 1:
						// [preds: 0 1] [succs: 2 1]
						// 159:
						// 160: if l4 != -1 goto <159>
						if (targetBlkIdx > j || targetBlkIdx <= i) {
							List args = new ArrayList();
							args.add(StringConstant.v(className));
							args.add(StringConstant.v(methodName));
							args.add(StringConstant.v(startIs));
							args.add(IntConstant.v(k));
							args.add(IntConstant.v(targetBlkIdx));
							Block instruBlk = blockLs.get(targetBlkIdx);
							InvokeExpr incExpr = Jimple.v()
									.newStaticInvokeExpr(flushES.makeRef(),
											args);
							InvokeStmt incStmt = Jimple.v().newInvokeStmt(
									incExpr);							
							insertUnitToHead(incStmt, instruBlk);
						}
						k++;
					}
					i = j;
				} else if (ut instanceof TableSwitchStmt
						|| ut instanceof LookupSwitchStmt) {
					Block block = blockLs.get(i);
					List<Block> sucBs = block.getSuccs();

					for (Block sucb : sucBs) {
						Block instruBlk = sucb;
						Unit uh = instruBlk.getHead();
						if (!(uh instanceof IdentityStmt)) {
							List args = new ArrayList();
							args.add(StringConstant.v(className));
							args.add(StringConstant.v(methodName));
							args.add(StringConstant.v(startIs));
							args.add(IntConstant.v(i));
							args.add(IntConstant.v(sucb.getIndexInMethod()));
							// here we don't differ the default case and other
							// case

							InvokeExpr incExpr = Jimple.v()
									.newStaticInvokeExpr(flushES.makeRef(),
											args);
							InvokeStmt incStmt = Jimple.v().newInvokeStmt(
									incExpr);

							instruBlk.insertBefore(incStmt, uh);
							uh.redirectJumpsToThisTo(incStmt);
						}
					}
				}
			}
		}
				
		while(blockIt.hasNext()) {
			Block block = (Block) blockIt.next();
			Unit uh = block.getHead();
			Unit ut = block.getTail();
			int blkIdx = block.getIndexInMethod();
			
			if(uh instanceof IdentityStmt && 
				!(ut instanceof IfStmt || ut instanceof TableSwitchStmt
					|| ut instanceof LookupSwitchStmt )) {
				List argxs = new ArrayList();
				argxs.add(StringConstant.v(className));
				argxs.add(StringConstant.v(methodName));
				argxs.add(StringConstant.v("-1>"+blkIdx));
				InvokeExpr insExpr = Jimple.v().newStaticInvokeExpr(add.makeRef(),
						argxs);
				InvokeStmt insStmt = Jimple.v().newInvokeStmt(insExpr);
				insertUnitToHead(insStmt, block);
			}
			//insert first mute predicate which will always be true meaning coverage this new code partition
			
			
			List args = new ArrayList();
			args.add(StringConstant.v(className));
			args.add(StringConstant.v(methodName));
			args.add(IntConstant.v(blkIdx));
			InvokeExpr incExpr1 = Jimple.v().newStaticInvokeExpr(
					increaseCounter.makeRef(), args);			
			InvokeStmt incStmt1 = Jimple.v().newInvokeStmt(incExpr1);
			insertUnitToHead(incStmt1, block);
			// must insirt in head but after identityStmt becuase if you
			// insirt before tail, when run this block, the stmt before the
			// counter stmt throw exception, then this counter will not be
			// recorded
		
		
			if (ut instanceof IfStmt || ut instanceof TableSwitchStmt
					|| ut instanceof LookupSwitchStmt) {
				InvokeExpr incExpr2 = Jimple.v().newStaticInvokeExpr(
						increaseES.makeRef(), args);
				InvokeStmt incStmt2 = Jimple.v().newInvokeStmt(incExpr2);
				block.insertBefore(incStmt2, ut);
				ut.redirectJumpsToThisTo(incStmt2);
			}
			//give every branch statement a ES record
		}
		
	}
	
	private void insertUnitToHead(Unit insertUnit, Block insrtBlock) {
		Unit uh = insrtBlock.getHead();
		
		if (!(uh instanceof IdentityStmt)) {
			insrtBlock.insertBefore(insertUnit, uh);
			uh.redirectJumpsToThisTo(insertUnit);
		} else {
			Unit tmp = uh;
			while (tmp != null && tmp instanceof IdentityStmt) {
				tmp = insrtBlock.getSuccOf(tmp);
			}
			if (tmp != null)
				insrtBlock.insertBefore(insertUnit, tmp);
			else
				insrtBlock.insertAfter(insertUnit, insrtBlock.getTail());
		}
	}
	
}